import torch
import math
import os


class CkptManager:
    def __init__(self, root, hist_cnt, best_mode, verbose=False):
        self.root = root
        os.makedirs(root, exist_ok=True)
        self.hist_cnt = hist_cnt
        self.hist_dict = {i: None for i in range(hist_cnt)}
        self.saved_ckpt_counter = 0
        self.verbose = verbose
        self.best_mode = best_mode  # None means not save best ckpt
        if best_mode:
            self.best_value = {"smaller": math.inf, "bigger": -math.inf}[best_mode]

    def save(self, save_this_dict, name, current_value=None, anchor=False):
        """save ckpt

        :param save_this_dict: dict to be saved
        :param name: name of the ckpt
        :param current_value: current value for best model metric
        :param anchor: boolean, if True will backup this ckpt by: old_name_anchor
        :returns: 

        """
        if anchor:
            anchor_path = f"{self.root}/{name}_anchor.pth"
            torch.save(save_this_dict, anchor_path)
            if self.verbose:
                print(f"ckpt {name} saved as anchor in {anchor_path}")
            return
        hist_id = self.saved_ckpt_counter % self.hist_cnt
        if self.hist_dict[hist_id] is not None:
            if os.path.exists(self.hist_dict[hist_id]):
                os.remove(self.hist_dict[hist_id])
        ckpt_path = f"{self.root}/{name}.pth"
        save_this_dict["ckpt_manager"] = self
        torch.save(save_this_dict, ckpt_path)
        self.hist_dict[hist_id] = ckpt_path
        if self.verbose:
            print(f"model saved in {ckpt_path}")
        if self.best_mode:
            if self.best_mode == "smaller" and current_value < self.best_value:
                torch.save(save_this_dict, f"{self.root}/best.pth")
                if self.verbose:
                    print(f"best model updated by {name}")
            elif self.best_mode == "bigger" and current_value > self.best_value:
                torch.save(save_this_dict, f"{self.root}/best.pth")
                if self.verbose:
                    print(f"best model updated by {name}")
            else:
                raise ValueError(f"best mode: {self.best_mode} is not supported!!!")
        self.saved_ckpt_counter += 1

    def init(
        self, new_root=None, new_hist_cnt=None, new_best_mode=None, new_verbose=None
    ):
        if new_root:
            self.root = new_root
        if new_hist_cnt:
            self.hist_cnt = new_hist_cnt
        self.hist_dict = {i: None for i in range(self.hist_cnt)}
        self.saved_ckpt_counter = 0
        if self.new_verbose:
            self.verbose = new_verbose
        if new_best_mode:
            self.best_mode = new_best_mode
            self.best_value = {
                "smaller": math.inf,
                "bigger": -math.inf,
                "latest": None,
            }[self.best_mode]

    def load(self, name, root=None):
        root = self.root if root is None else root
        return torch.load(f"{root}/{name}", map_location="cpu")
